package com.edu.ocrdemo.service;

import com.edu.ocrdemo.common.TaskStatus;
import com.edu.ocrdemo.entity.Task;
import com.edu.ocrdemo.repository.TaskRepository;
import lombok.AllArgsConstructor;
import net.sourceforge.tess4j.Tesseract;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;

import java.awt.*;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.awt.image.ConvolveOp;
import java.awt.image.Kernel;
import java.io.ByteArrayInputStream;
import java.io.InputStream;

import java.time.LocalDateTime;

/**
 * @Author tfan
 * @Description
 * @Date 2024/12/31 17:17
 **/
@Service
@AllArgsConstructor
public class FileProcessor {

    private final Tesseract tesseract;
    private final TaskRepository taskRepository;




    @Async
    public void processFile(String taskId, MultipartFile[] files) {
        Task task = taskRepository.findById(taskId);
        if (task == null) {
            return;
        }

        task.setStatus(TaskStatus.PROCESSING);
        StringBuilder ocrStr = new StringBuilder();
        try {
            for (MultipartFile file : files) {
                // 转换
                InputStream sbs = new ByteArrayInputStream(file.getBytes());
                BufferedImage bufferedImage = ImageIO.read(sbs);
                BufferedImage preprocessedImage = preprocessImage(bufferedImage);
                // 对图片进行文字识别
                ocrStr.append(tesseract.doOCR(preprocessedImage));
            }
            // 处理成功后更新任务状态
            task.setStatus(TaskStatus.COMPLETED);
            task.setResult(ocrStr.toString());
            task.setEndTime(LocalDateTime.now());
        } catch (Exception e) {
            task.setStatus(TaskStatus.FAILED);
            task.setResult("Error processing file: " + e.getMessage());
        }
    }

    public BufferedImage preprocessImage(BufferedImage original) {
        // 1. 调整图像大小
        BufferedImage resized = resize(original, original.getWidth() * 2, original.getHeight() * 2);

        // 2. 转换为灰度图
        BufferedImage grayscale = toGrayscale(resized);

        // 3. 降噪
        BufferedImage denoised = denoise(grayscale);

        // 4. 二值化
        BufferedImage binary = toBinary(denoised);

        // 5. 锐化
        BufferedImage sharpened = sharpen(binary);

        return sharpened;
    }

    private BufferedImage resize(BufferedImage original, int width, int height) {
        Image tmp = original.getScaledInstance(width, height, Image.SCALE_SMOOTH);
        BufferedImage resized = new BufferedImage(width, height, BufferedImage.TYPE_INT_RGB);
        Graphics2D g2d = resized.createGraphics();
        g2d.drawImage(tmp, 0, 0, null);
        g2d.dispose();
        return resized;
    }

    private BufferedImage toGrayscale(BufferedImage original) {
        BufferedImage grayscale = new BufferedImage(
                original.getWidth(),
                original.getHeight(),
                BufferedImage.TYPE_BYTE_GRAY);
        Graphics2D g2d = grayscale.createGraphics();
        g2d.drawImage(original, 0, 0, null);
        g2d.dispose();
        return grayscale;
    }

    private BufferedImage denoise(BufferedImage original) {
        // 使用中值滤波进行降噪
        int width = original.getWidth();
        int height = original.getHeight();
        BufferedImage result = new BufferedImage(width, height, original.getType());

        for (int y = 1; y < height - 1; y++) {
            for (int x = 1; x < width - 1; x++) {
                int[] pixels = new int[9];
                int idx = 0;

                // 获取3x3窗口内的像素值
                for (int i = -1; i <= 1; i++) {
                    for (int j = -1; j <= 1; j++) {
                        pixels[idx++] = new Color(original.getRGB(x + j, y + i)).getRed();
                    }
                }

                // 排序找中值
                java.util.Arrays.sort(pixels);
                int median = pixels[4];

                result.setRGB(x, y, new Color(median, median, median).getRGB());
            }
        }
        return result;
    }

    private BufferedImage toBinary(BufferedImage original) {
        // 使用OTSU算法进行自适应阈值计算
        int width = original.getWidth();
        int height = original.getHeight();
        BufferedImage result = new BufferedImage(width, height, BufferedImage.TYPE_BYTE_BINARY);

        // 计算直方图
        int[] histogram = new int[256];
        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width; x++) {
                int gray = new Color(original.getRGB(x, y)).getRed();
                histogram[gray]++;
            }
        }

        // OTSU算法计算阈值
        int threshold = getOtsuThreshold(histogram, width * height);

        // 应用阈值
        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width; x++) {
                int gray = new Color(original.getRGB(x, y)).getRed();
                int newPixel = gray > threshold ? 255 : 0;
                result.setRGB(x, y, new Color(newPixel, newPixel, newPixel).getRGB());
            }
        }
        return result;
    }

    private int getOtsuThreshold(int[] histogram, int total) {
        float sum = 0;
        for (int i = 0; i < 256; i++) {
            sum += i * histogram[i];
        }

        float sumB = 0;
        int wB = 0;
        int wF;
        float varMax = 0;
        int threshold = 0;

        for (int i = 0; i < 256; i++) {
            wB += histogram[i];
            if (wB == 0) continue;

            wF = total - wB;
            if (wF == 0) break;

            sumB += i * histogram[i];
            float mB = sumB / wB;
            float mF = (sum - sumB) / wF;

            float varBetween = (float) wB * (float) wF * (mB - mF) * (mB - mF);

            if (varBetween > varMax) {
                varMax = varBetween;
                threshold = i;
            }
        }
        return threshold;
    }

    private BufferedImage sharpen(BufferedImage original) {
        float[] sharpenMatrix = {
                0.0f, -1.0f, 0.0f,
                -1.0f, 5.0f, -1.0f,
                0.0f, -1.0f, 0.0f
        };

        Kernel kernel = new Kernel(3, 3, sharpenMatrix);
        ConvolveOp op = new ConvolveOp(kernel, ConvolveOp.EDGE_NO_OP, null);
        return op.filter(original, null);
    }}

